#pragma once

#include "viscous_force.h"

namespace SPH
{
namespace fluid_dynamics
{
//=================================================================================================//
template <typename ViscosityType, class KernelCorrectionType,
          template <typename...> class RelationType, typename... Parameters>
template <class BaseRelationType>
ViscousForceCK<Base, ViscosityType, KernelCorrectionType, RelationType<Parameters...>>::
    ViscousForceCK(BaseRelationType &base_relation)
    : Interaction<RelationType<Parameters...>>(base_relation),
      viscosity_model_(DynamicCast<ViscosityType>(this, this->particles_->getBaseMaterial())),
      kernel_correction_(this->particles_),
      dv_vel_(this->particles_->template getVariableByName<Vecd>("Velocity")),
      dv_viscous_force_(this->particles_->template registerStateVariable<Vecd>("ViscousForce")),
      smoothing_length_sq_(pow(this->getSPHAdaptation().ReferenceSmoothingLength(), 2)) {}
//=================================================================================================//
template <typename ViscosityType, class KernelCorrectionType, typename... Parameters>
ViscousForceCK<Inner<WithUpdate, ViscosityType, KernelCorrectionType, Parameters...>>::
    ViscousForceCK(Inner<Parameters...> &inner_relation)
    : BaseViscousForceType(inner_relation),
      ForcePriorCK(this->particles_, this->dv_viscous_force_) {}
//=================================================================================================//
template <typename ViscosityType, class KernelCorrectionType, typename... Parameters>
template <class ExecutionPolicy, class EncloserType>
ViscousForceCK<Inner<WithUpdate, ViscosityType, KernelCorrectionType, Parameters...>>::
    InteractKernel::InteractKernel(const ExecutionPolicy &ex_policy, EncloserType &encloser)
    : BaseViscousForceType::InteractKernel(ex_policy, encloser),
      inter_particle_viscosity_(
          encloser.viscosity_model_.getInterParticleViscosity(ex_policy, encloser.viscosity_model_)),
      correction_(ex_policy, encloser.kernel_correction_),
      Vol_(encloser.dv_Vol_->DelegatedData(ex_policy)),
      vel_(encloser.dv_vel_->DelegatedData(ex_policy)),
      viscous_force_(encloser.dv_viscous_force_->DelegatedData(ex_policy)),
      smoothing_length_sq_(encloser.smoothing_length_sq_) {}
//=================================================================================================//
template <typename ViscosityType, class KernelCorrectionType, typename... Parameters>
void ViscousForceCK<Inner<WithUpdate, ViscosityType, KernelCorrectionType, Parameters...>>::
    InteractKernel::interact(size_t index_i, Real dt)
{
    Vecd force = Vecd::Zero();
    for (UnsignedInt n = this->FirstNeighbor(index_i); n != this->LastNeighbor(index_i); ++n)
    {
        UnsignedInt index_j = this->neighbor_index_[n];
        Vecd e_ij = this->e_ij(index_i, index_j);
        Real dW_ijV_j = this->dW_ij(index_i, index_j) * Vol_[index_j];
        Vecd vec_r_ij = this->vec_r_ij(index_i, index_j);

        Vecd vel_derivative = (vel_[index_i] - vel_[index_j]) /
                              (vec_r_ij.squaredNorm() + 0.01 * smoothing_length_sq_);

        force += vec_r_ij.dot((correction_(index_i) + correction_(index_j)) * e_ij) *
                 inter_particle_viscosity_(index_i, index_j) * vel_derivative * dW_ijV_j;
    }
    viscous_force_[index_i] = force * Vol_[index_i];
}
//=================================================================================================//
template <typename ViscosityType, class KernelCorrectionType, typename... Parameters>
ViscousForceCK<Contact<Wall, ViscosityType, KernelCorrectionType, Parameters...>>::
    ViscousForceCK(Contact<Parameters...> &contact_relation)
    : BaseViscousForceType(contact_relation), Interaction<Wall>(contact_relation) {}
//=================================================================================================//
template <typename ViscosityType, class KernelCorrectionType, typename... Parameters>
template <class ExecutionPolicy, class EncloserType>
ViscousForceCK<Contact<Wall, ViscosityType, KernelCorrectionType, Parameters...>>::InteractKernel::
    InteractKernel(const ExecutionPolicy &ex_policy, EncloserType &encloser, size_t contact_index)
    : BaseViscousForceType::InteractKernel(ex_policy, encloser, contact_index),
      one_side_viscosity_(encloser.viscosity_model_.getOneSideViscosity(ex_policy)),
      correction_(ex_policy, encloser.kernel_correction_),
      Vol_(encloser.dv_Vol_->DelegatedData(ex_policy)),
      contact_Vol_(encloser.dv_contact_Vol_[contact_index]->DelegatedData(ex_policy)),
      vel_(encloser.dv_vel_->DelegatedData(ex_policy)),
      viscous_force_(encloser.dv_viscous_force_->DelegatedData(ex_policy)),
      wall_vel_ave_(encloser.dv_wall_vel_ave_[contact_index]->DelegatedData(ex_policy)),
      smoothing_length_sq_(encloser.smoothing_length_sq_) {}
//=================================================================================================//
template <typename ViscosityType, class KernelCorrectionType, typename... Parameters>
void ViscousForceCK<Contact<Wall, ViscosityType, KernelCorrectionType, Parameters...>>::
    InteractKernel::interact(size_t index_i, Real dt)
{
    Vecd force = Vecd::Zero();
    for (UnsignedInt n = this->FirstNeighbor(index_i); n != this->LastNeighbor(index_i); ++n)
    {
        UnsignedInt index_j = this->neighbor_index_[n];
        Vecd e_ij = this->e_ij(index_i, index_j);
        Real dW_ijV_j = this->dW_ij(index_i, index_j) * contact_Vol_[index_j];
        Vecd vec_r_ij = this->vec_r_ij(index_i, index_j);

        Vecd vel_derivative = 2.0 * (vel_[index_i] - wall_vel_ave_[index_j]) /
                              (vec_r_ij.squaredNorm() + 0.01 * smoothing_length_sq_);

        force += 2.0 * vec_r_ij.dot(correction_(index_i) * e_ij) *
                 one_side_viscosity_(index_i) * vel_derivative * dW_ijV_j;
    }
    viscous_force_[index_i] += force * Vol_[index_i];
}
//=================================================================================================//
} // namespace fluid_dynamics
} // namespace SPH
